#导入必要的包
import cv2
import gradio as gr
import numpy as np
from sklearn.cluster import KMeans  # 实现K-means聚类

# 将图像的像素颜色数据聚类成k=5个中心，并返回这些中心的颜色。
def kmeans_clustering_sklearn(image, k=5):
    # 1. 准备好 RGB 三列的像素颜色矩阵
    pixel = image.reshape(-1, 3)  
    # 2. 通过K-means聚类对图像进行聚类
    kmeans = KMeans(n_clusters=k, random_state=0, n_init=10) 
    kmeans.fit(pixel) 
    colors = kmeans.cluster_centers_   # 获取聚类中心的像素颜色
    # 3. 将颜色值四舍五入到最接近的整数
    colors = colors.round().astype(int)   
    return colors

# 定义 gradio 应用函数
def gradio_app(image):
    # 提取图像的主要颜色
    colors = kmeans_clustering_sklearn(image)
    # 生成十六进制颜色代码 
    hexadecimal_colors = ["#%02x%02x%02x" % (color[0], color[1], color[2]) for color in colors]
    # 创建一个空的数组，用于存放调色板图像
    back = np.zeros((100, 100 * len(colors), 3), dtype=np.uint8)
    for i, color in enumerate(colors):
        back[:, i * 100:(i + 1) * 100] = color
    
    # 返回颜色调色板图像和颜色调色板的十六进制表示
    return back, "\n".join(hexadecimal_colors)
  
# 创建 gradio 用户界面  
iface = gr.Interface(  
    fn=gradio_app, # 指定 gradio 应用函数  
    inputs=gr.components.Image(), # 定义输入类型为图像  
      
    # 定义输出类型为图像、文本和颜色列表  
    outputs=[
        gr.components.Image(type="numpy",show_label=False,width=300),
        gr.components.Text(label="色彩十六进制码")
    ], 
    title="KMeans调色板提取器", # 定义应用标题  
    description="上传图像并获取其调色板。" # 定义应用描述  
)  
  
iface.launch(share=True) # 启动 gradio 应用  
